import requests
import json
from pathlib import Path
import pandas as pd
import random
import time
import os
import re  # 用于解析模型B的输出

# --- 全局变量 ---
CONFIG = {}  # 将由配置加载函数填充
API_KEY_MODEL_B = None  # 将在加载配置后设置


# --- 配置加载函数 ---
def load_validation_configuration(config_filename="config_validation.json"):
    global CONFIG, API_KEY_MODEL_B
    try:
        script_dir = Path(__file__).resolve().parent
        config_path = script_dir / config_filename
        with open(config_path, 'r', encoding='utf-8') as f:
            CONFIG = json.load(f)
        print(f"✅ 校验配置文件 '{config_filename}' 加载成功。")

        # 从环境变量加载模型B的API Key (环境变量名称来自配置文件)
        api_key_env_var = CONFIG.get("model_b_api_settings", {}).get("api_key_env_var")
        if not api_key_env_var:
            print("错误：校验配置文件中未指定 'model_b_api_settings.api_key_env_var'。")
            return False

        API_KEY_MODEL_B = os.getenv(api_key_env_var)
        if not API_KEY_MODEL_B:
            print(f"错误：找不到模型B的 API Key。请设置 '{api_key_env_var}' 环境变量。")
            return False
        print(f"✅ 模型B的API Key（环境变量 {api_key_env_var}）加载成功。")
        return True

    except FileNotFoundError:
        print(f"错误：找不到校验配置文件 {config_path}")
    except json.JSONDecodeError:
        print(f"错误：校验配置文件 {config_path} 格式有误。")
    except Exception as e:
        print(f"加载校验配置文件 {config_path} 时出错: {e}")
    return False


# --- 模型B API调用函数 (示例为DashScope风格) ---
def call_model_b_api(system_prompt: str, user_prompt_for_validation: str) -> dict:
    global API_KEY_MODEL_B, CONFIG

    if 'model_b_api_settings' not in CONFIG:
        print("错误：配置中缺少 'model_b_api_settings'。")
        return {"text": "错误：模型B API配置缺失", "usage": {}, "request_id": ""}

    api_conf = CONFIG['model_b_api_settings']

    url = f"https://{api_conf['api_host']}{api_conf['api_path_template']}"
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {API_KEY_MODEL_B}'
    }
    payload = {
        "model": api_conf['model_id_for_payload'],
        "input": {
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt_for_validation}
            ]
        },
        "parameters": api_conf.get("default_model_params", {})
    }

    try:
        response = requests.post(
            url,
            headers=headers,
            json=payload,
            timeout=CONFIG['script_parameters'].get('api_timeout_seconds', 120)
        )
        response.raise_for_status()
        response_data = response.json()
        text_output = response_data.get("output", {}).get("text", "")
        usage_data = response_data.get("usage", {})

        return {
            "text": text_output,
            "usage": {
                "input_tokens": usage_data.get("input_tokens"),
                "output_tokens": usage_data.get("output_tokens"),
                "total_tokens": usage_data.get("input_tokens", 0) + usage_data.get("output_tokens", 0)
            },
            "request_id": response_data.get("request_id", "")
        }
    except requests.exceptions.Timeout:
        print("    -> 错误：模型B API 请求超时。")
        return {"text": "错误：请求超时", "usage": {}, "request_id": ""}
    except requests.exceptions.RequestException as e:
        error_message = f"错误：模型B API 请求失败: {e}"
        # ... (错误处理同之前的脚本) ...
        print(f"    -> {error_message}")
        return {"text": f"错误：{str(e)}", "usage": {}, "request_id": ""}


# --- 解析模型B的输出 ---
# --- 解析模型B的输出 (增强版) ---
def parse_model_b_response(response_text: str):
    """
    解析模型B的响应文本，提取校验码 (0或1) 和理由。
    更灵活地处理可能的输出格式。
    """
    validation_code = None
    reason_or_suggestion = response_text.strip() # 默认为全部剥离首尾空格的响应文本

    # 模式1: 尝试匹配 "判断：0" 或 "判断：1" (忽略大小写，处理可选的**和空格)
    # 例如: "判断：0", "**判断： 0**", "判断:1"
    match_keyword = re.search(r"(?:\*\*| JUDGEMENT:\s*)?判断\s*(?:：|:)\s*([01])", reason_or_suggestion, re.IGNORECASE) # 添加了 JUDGEMENT
    if match_keyword:
        validation_code = match_keyword.group(1)
        # 获取 "判断：0" 模式之后的内容作为理由
        split_point = match_keyword.end(0)
        reason_or_suggestion = reason_or_suggestion[split_point:].strip()
        # 进一步清理理由开头可能存在的标点或特定词
        reason_or_suggestion = re.sub(r"^\s*(?:理由\s*(?:：|:)|，|。|：|,|\.|:)\s*", "", reason_or_suggestion, flags=re.IGNORECASE).strip()
        return validation_code, reason_or_suggestion

    # 模式2: 如果上面没有匹配到，再尝试直接匹配开头的 "0" 或 "1"
    match_simple = re.match(r"^\s*([01])(.*)", reason_or_suggestion, re.DOTALL)
    if match_simple:
        validation_code = match_simple.group(1)
        reason_or_suggestion = match_simple.group(2).strip()
        # 清理理由开头可能存在的标点
        reason_or_suggestion = re.sub(r"^\s*(?:理由\s*(?:：|:)|，|。|：|,|\.|:)\s*", "", reason_or_suggestion, flags=re.IGNORECASE).strip()
        return validation_code, reason_or_suggestion

    # 回退逻辑: 如果以上模式都未匹配，检查文本中是否唯一包含 "0" 或 "1"
    # 这个回退逻辑的可靠性较低，因为 "0" 或 "1" 可能出现在理由文本中
    # 但作为最后尝试，如果模型完全不按格式输出
    # if validation_code is None: # 只有在上面两种模式都失败时才执行
    #     contains_0 = "0" in reason_or_suggestion
    #     contains_1 = "1" in reason_or_suggestion
    #     if contains_0 and not contains_1:
    #         validation_code = "0"
    #         # 此时，reason_or_suggestion 仍然是完整文本，因为不确定如何精确分割
    #     elif contains_1 and not contains_0:
    #         validation_code = "1"
            # 此时，reason_or_suggestion 仍然是完整文本

    # 如果经过所有尝试，validation_code 仍为 None，说明模型输出格式严重不符
    # reason_or_suggestion 此时是原始的、剥离了首尾空格的响应文本
    return validation_code, reason_or_suggestion


# --- Excel数据加载函数 (与之前类似，但读取特定列) ---
def load_excel_for_validation(file_settings_config):
    try:
        file_path = Path(__file__).resolve().parent / file_settings_config["input_excel_file_from_model_a"]
        df = pd.read_excel(file_path)  # 默认读取第一个sheet，如果需要指定sheet，需要修改

        # 检查所需列是否存在
        required_cols = [
            file_settings_config["original_text_column_name"],
            file_settings_config["model_a_label_column_name"]
        ]
        missing_cols = [col for col in required_cols if col not in df.columns]
        if missing_cols:
            print(f"错误：输入Excel文件 '{file_path}' 缺少以下必需列: {', '.join(missing_cols)}")
            return None

        print(f"✅ 成功从 '{file_path}' 加载数据进行校验 (共 {len(df)} 行)。")
        return df
    except FileNotFoundError:
        print(f"错误：找不到输入Excel文件 {file_path}")
    except Exception as e:
        print(f"加载输入Excel文件 {file_path} 时出错: {e}")
    return None


# --- 保存函数 (与之前类似) ---
def save_validation_data_to_excel(data_list, output_file_key, chunk_save=False):
    global CONFIG  # 确保CONFIG已加载
    if not data_list:
        print("没有校验数据需要保存。")
        return

    df_new = pd.DataFrame(data_list)
    script_dir = Path(__file__).resolve().parent
    file_name = CONFIG["validation_file_settings"][output_file_key]  # 从validation_file_settings获取
    file_path = script_dir / file_name

    mode = 'a' if chunk_save and file_path.exists() else 'w'
    header = not (chunk_save and file_path.exists())  # 如果是追加模式且文件存在，则不写header

    try:
        if mode == 'a':
            with pd.ExcelWriter(file_path, engine='openpyxl', mode='a', if_sheet_exists='overlay') as writer:
                # 获取最后一个sheet的名称，如果文件是新创建的，则会出错，需要处理
                try:
                    df_existing = pd.read_excel(file_path, sheet_name=None)  # 读取所有sheet
                    last_sheet_name = list(df_existing.keys())[-1]
                    startrow = writer.sheets[last_sheet_name].max_row
                    df_new.to_excel(writer, sheet_name=last_sheet_name, startrow=startrow, index=False, header=False)
                except Exception:  # 文件可能是空的或格式问题，直接覆盖写入新sheet
                    df_new.to_excel(writer, sheet_name="ValidationResults", index=False, header=True)

            print(f"已追加 {len(data_list)} 条校验数据到 {file_name}")
        else:  # 'w' mode
            df_new.to_excel(file_path, index=False, engine='openpyxl', sheet_name="ValidationResults")
            action = "创建并写入" if chunk_save else "成功将"  # chunk_save=False时，总是覆盖创建
            print(f"✅ {action} {len(data_list)} 条校验数据保存到 {file_name}")

    except Exception as e:
        print(f"保存或追加到 Excel 文件 {file_name} 时出错: {e}. 尝试直接写入新文件/覆盖...")
        try:  # 备用：直接覆盖写入
            df_new.to_excel(file_path, index=False, engine='openpyxl', sheet_name="ValidationResults")
            print(f"已将 {len(data_list)} 条校验数据覆盖写入到 {file_name}")
        except Exception as e_final:
            print(f"最终写入 Excel 文件 {file_name} 失败: {e_final}")


# --- 主逻辑 ---
if __name__ == "__main__":
    if not load_validation_configuration():
        exit()

    file_settings = CONFIG["validation_file_settings"]
    prompt_settings = CONFIG["validation_prompt_settings"]
    script_params = CONFIG["script_parameters"]

    system_prompt_b = prompt_settings["system_prompt_for_model_b"]
    user_prompt_template_b = prompt_settings["cross_validation_user_prompt_template"]

    # 加载模型A输出的Excel数据
    df_to_validate = load_excel_for_validation(file_settings)
    if df_to_validate is None:
        print("错误：无法加载待校验数据，程序退出。")
        exit()

    original_text_col = file_settings["original_text_column_name"]
    model_a_label_col = file_settings["model_a_label_column_name"]

    # 新增的列名
    model_b_code_col = file_settings["model_b_validation_code_column"]
    model_b_reason_col = file_settings["model_b_reason_column"]

    validation_results_chunk = []
    total_validated_count = 0
    total_model_b_requests = 0

    print("\n--- 开始进行标签校验 (调用模型B API) ---")

    for index, row in df_to_validate.iterrows():
        original_text = str(row.get(original_text_col, ""))
        model_a_labels = str(row.get(model_a_label_col, ""))

        excel_row_num = index + 2  # Excel行号（假设第一行为标题）
        print(
            f"  [请求 {total_model_b_requests + 1}] 校验第 {excel_row_num} 行 - 模型A标签: '{model_a_labels[:50]}...'")

        if not original_text.strip() or not model_a_labels.strip():
            print("    -> 原始文本或模型A标签为空，跳过校验。")
            validation_code = "N/A"
            reason_or_suggestion = "输入为空或不完整"
            api_response_b = {}
        else:
            # 格式化校验Prompt
            validation_prompt = user_prompt_template_b.format(
                original_text=original_text,
                model_a_labels=model_a_labels
            )
            total_model_b_requests += 1
            api_response_b = call_model_b_api(system_prompt_b, validation_prompt)
            model_b_raw_output = api_response_b.get("text", "错误：未能获取模型B响应")
            validation_code, reason_or_suggestion = parse_model_b_response(model_b_raw_output)

        # 准备要保存的数据行
        output_row = row.to_dict()  # 复制原始行数据
        output_row[model_b_code_col] = validation_code
        output_row[model_b_reason_col] = reason_or_suggestion
        # 可以选择性地添加API B的调用信息
        output_row["模型B_API请求ID"] = api_response_b.get("request_id", "")
        output_row["模型B_输入Tokens"] = api_response_b.get("usage", {}).get("input_tokens")
        output_row["模型B_输出Tokens"] = api_response_b.get("usage", {}).get("output_tokens")
        output_row["校验时间"] = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())

        validation_results_chunk.append(output_row)
        total_validated_count += 1

        print(
            f"    -> 模型B判断: '{validation_code}', 理由/建议: '{str(reason_or_suggestion)[:50]}...'，当前批次 {len(validation_results_chunk)} 条。")

        if len(validation_results_chunk) >= script_params["save_interval"]:
            save_validation_data_to_excel(validation_results_chunk, "output_validation_excel_file", chunk_save=True)
            validation_results_chunk = []

        time.sleep(random.uniform(
            script_params["request_delay_min_seconds"],
            script_params["request_delay_max_seconds"]
        ))

    print("\n--- 所有数据校验完成 ---")
    if validation_results_chunk:
        print(f"正在保存最后剩余的 {len(validation_results_chunk)} 条校验数据...")
        save_validation_data_to_excel(validation_results_chunk, "output_validation_excel_file", chunk_save=True)
    else:
        print("没有剩余校验数据需要保存。")

    output_filename_val = CONFIG['validation_file_settings']['output_validation_excel_file']
    print(f"\n总计校验了 {total_validated_count} 行数据，向模型B发起了 {total_model_b_requests} 次API请求。")
    print(f"✅ 所有校验数据已尝试保存到 Excel 文件中 '{output_filename_val}'。")